import os
from select import select
import sys
import cv2
import datetime


videoPath = r'./preview.mp4'

jpgSavePath = r'./darknet/darknet-master/build/darknet/x64/data/VOCdevkit2007/VOC2007/JPEGImages/'
# 在路径前边加了个当前目录 './'， 是测试时使用的，可以去掉

# 文件夹必须存在才能保存成功
if not os.path.exists(jpgSavePath):
    os.makedirs(jpgSavePath)

# 视频帧截取
cap = cv2.VideoCapture(videoPath)

# 摄像头画面截取
#cap = cv2.VideoCapture(0)

if cap.isOpened():
	# 判断是否正常打开
    cap_flag, frame = cap.read()
    print('VideoCapture Status: ' + repr(cap_flag))
else:
    cap_flag = False
    print('VideoCapture Status: ' + repr(cap_flag))
    sys.exit(0)

# 视频信息
frame_FORMAT = cap.get(cv2.CAP_PROP_FORMAT) # == cap.get(8)
print('视频格式: ', frame_FORMAT)
frame_COUNT = cap.get(cv2.CAP_PROP_FRAME_COUNT) # == cap.get(7)
print('视频总帧数: ', frame_COUNT)
frame_BITRATE = cap.get(cv2.CAP_PROP_BITRATE) # == cap.get(47)
print('数据速率: ', frame_BITRATE)
frame_FPS = cap.get(cv2.CAP_PROP_FPS) # == cap.get(5)
print('视频帧率: ' + repr(frame_FPS))
frame_FH = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) # == cap.get(4)
print('视频帧高: ' + repr(frame_FH))
frame_FW = cap.get(cv2.CAP_PROP_FRAME_WIDTH) # == cap.get(3)
print('视频帧宽: ' + repr(frame_FW))
m, s = divmod(frame_COUNT/frame_FPS, 60)
h, m = divmod(m, 60)
print('视频时长: ' + "%d:%02d:%02d" % (h, m, s) )
print('\n')

sel = input('选择截取视频帧的模式（t - 按固定时间间隔截取；f - 按固定帧间隔截取）:  ')
if(sel == 't' or sel == 'T'): 
    # 截取视频帧的时间间隔
    timeRate = input('请输入截取视频帧的时间间隔（单位：秒，按回车默认5秒）:  ')
    if(timeRate == ''):
        timeRate = 5  
    else:
        timeRate = int(timeRate)
    frameRate = int(frame_FPS) * int(timeRate) # 转换为帧间隔
elif(sel == 'f' or sel == 'F'):
    # 截取视频帧的帧间隔
    frameRate = input('请输入截取视频帧的帧间隔（单位：帧，按回车默认60帧）:  ')
    if(frameRate == ''):
        frameRate = 60  
    else:
        frameRate = int(frameRate)
else:
    cap.release()
    sys.exit(0)

print('截取视频帧的帧间隔: '+ repr(frameRate))

cnt = 1 # 从第一帧开始计数
pic_index = 0 # 图片序列号
while(cap_flag):
    ret, frame = cap.read()

    if(ret):
        if(cnt % frameRate == 0): # 每固定帧间隔截取图片保存
            jpgSaveName = jpgSavePath + '%06d'%pic_index + '.jpg' # 图片保存名字
            sret = cv2.imwrite(jpgSaveName, frame)
            if(sret):
                print('第'+ repr(cnt) +'帧 截取成功！'+ str(datetime.datetime.now()) + ' ' + jpgSaveName)
                pic_index+=1
            else:
                print('第'+ repr(cnt) +'帧 截取失败！')
        cnt+=1
        cv2.imshow('preview',frame) # 视频预览
        # cv2.waitKey(1)
        cv2.waitKey(1000//frameRate) # 原视频速率
    else: # 最后一帧为空帧
        print('从' + repr(cnt) +'帧中一共截取了'+ repr(pic_index) + '张图片。')
        break

cap.release()
